`include "defines.v"

module pipeline_regs (
    input  wire                         clock,
    input  wire                         reset,
    input  wire                         bubble,
    input  wire                         flush,
    input  wire [ 3: 0]                 stall,

    //IF->ID
    input  wire                         o_pc_valid_if,
    input  wire [`BUSLEN-1:0]           o_pc_if,
    input  wire [`INSTLEN-1:0]          o_instruction_if,

    output reg                          i_pc_valid_id,
    output reg  [`BUSLEN-1:0]           i_pc_id,
    output reg  [`INSTLEN-1:0]          i_instruction_id,

    //ID->EX
    input  wire [`XLEN-1:0]             o_regbus_A_id,
    input  wire [`XLEN-1:0]             o_regbus_B_id,
    input  wire [`XLEN-1:0]             o_imm_id,
    input  wire [ 1: 0]                 o_ALU_oprend_A_src_id,
    input  wire [ 1: 0]                 o_ALU_oprend_B_src_id,
    input  wire [ 2: 0]                 o_ALU_op_id,
    input  wire [ 2: 0]                 o_ALU_out_ext_type_id,
    input  wire                         o_adder_src_id,
    input  wire [`BUSLEN-1:0]           o_pc_id,
    input  wire                         o_pc_valid_id,

    input  wire                         o_mem_r_en_id,
    input  wire [ 2: 0]                 o_mem_r_ext_type_id,
    input  wire                         o_mem_w_en_id,
    input  wire [ 1: 0]                 o_mem_w_size_id,
    input  wire                         o_inst_jump_id,
    input  wire [ 2: 0]                 o_inst_branch_type_id,
    input  wire                         o_is_system_inst_id,
    input  wire                         o_inst_ecall_id,
    input  wire                         o_inst_mret_id,

    input  wire                         o_rd_w_en_id,
    input  wire                         o_rd_w_src_id,
    input  wire [`REGFILE_ADDR_LEN-1:0] o_rd_w_addr_id,
    input  wire                         o_csr_w_en_id,
    input  wire [`CSR_ADDR_LEN-1:0]     o_csr_w_addr_id,

    output reg  [`XLEN-1:0]             i_regbus_A_ex,
    output reg  [`XLEN-1:0]             i_regbus_B_ex,
    output reg  [`XLEN-1:0]             i_imm_ex,
    output reg  [ 1: 0]                 i_ALU_oprend_A_src_ex,
    output reg  [ 1: 0]                 i_ALU_oprend_B_src_ex,
    output reg  [ 2: 0]                 i_ALU_op_ex,
    output reg  [ 2: 0]                 i_ALU_out_ext_type_ex,
    output reg                          i_adder_src_ex,
    output reg  [`BUSLEN-1:0]           i_pc_ex,
    output reg                          i_pc_valid_ex,

    output reg                          i_mem_r_en_ex,
    output reg  [ 2: 0]                 i_mem_r_ext_type_ex,
    output reg                          i_mem_w_en_ex,
    output reg  [ 1: 0]                 i_mem_w_size_ex,
    output reg                          i_inst_jump_ex,
    output reg  [ 2: 0]                 i_inst_branch_type_ex,
    output reg                          i_is_system_inst_ex,
    output reg                          i_inst_ecall_ex,
    output reg                          i_inst_mret_ex,

    output reg                          i_rd_w_en_ex,
    output reg                          i_rd_w_src_ex,
    output reg  [`REGFILE_ADDR_LEN-1:0] i_rd_w_addr_ex,
    output reg                          i_csr_w_en_ex,
    output reg  [`CSR_ADDR_LEN-1:0]     i_csr_w_addr_ex,

    //EX->MEM
    input  wire [`BUSLEN-1:0]           o_pc_ex,

    input  wire [`BUSLEN-1:0]           o_ALU_out_ex,
    input  wire [ 3: 0]                 o_psw_flags_ex,

    input  wire                         o_mem_r_en_ex,
    input  wire [ 2: 0]                 o_mem_r_ext_type_ex,
    input  wire                         o_mem_w_en_ex,
    input  wire [ 1: 0]                 o_mem_w_size_ex,
    input  wire [`BUSLEN-1:0]           o_regbus_B_ex,

    input  wire [`REGFILE_ADDR_LEN-1:0] o_rd_w_addr_ex,
    input  wire                         o_rd_w_en_ex,
    input  wire                         o_rd_w_src_ex,
    input  wire                         o_csr_w_en_ex,
    input  wire [`CSR_ADDR_LEN-1:0]     o_csr_w_addr_ex,

    input  wire                         o_inst_jump_ex,
    input  wire [ 2: 0]                 o_inst_branch_type_ex,

    input  wire                         o_is_system_inst_ex,
    input  wire                         o_inst_ecall_ex,
    input  wire                         o_inst_mret_ex,

    output reg  [`BUSLEN-1:0]           i_branch_target_pc_mem,

    output reg  [`BUSLEN-1:0]           i_ALU_out_mem,
    output reg  [ 3: 0]                 i_psw_flags_mem,

    output reg                          i_mem_r_en_mem,
    output reg  [ 2: 0]                 i_mem_r_ext_type_mem,
    output reg                          i_mem_w_en_mem,
    output reg  [ 1: 0]                 i_mem_w_size_mem,
    output reg  [`BUSLEN-1:0]           i_regbus_B_mem,

    output reg  [`REGFILE_ADDR_LEN-1:0] i_rd_w_addr_mem,
    output reg                          i_rd_w_en_mem,
    output reg                          i_rd_w_src_mem,
    output reg                          i_csr_w_en_mem,
    output reg  [`CSR_ADDR_LEN-1:0]     i_csr_w_addr_mem,

    output reg                          i_inst_jump_mem,
    output reg  [ 2: 0]                 i_inst_branch_type_mem,

    output reg                          i_is_system_inst_mem,
    output reg                          i_inst_ecall_mem,
    output reg                          i_inst_mret_mem,
    
    //MEM->WB
    input  wire [`REGFILE_ADDR_LEN-1:0] o_rd_w_addr_mem,
    input  wire                         o_rd_w_en_mem,
    input  wire                         o_rd_w_src_mem,
    input  wire                         o_csr_w_en_mem,
    input  wire [`CSR_ADDR_LEN-1:0]     o_csr_w_addr_mem,
    input  wire [ 2: 0]                 o_mem_r_ext_type_mem,
    input  wire [`BUSLEN-1:0]           o_regbus_C_mem,
    input  wire [`BUSLEN-1:0]           o_regbus_D_mem,

    output reg  [`REGFILE_ADDR_LEN-1:0] i_rd_w_addr_wb,
    output reg                          i_rd_w_en_wb,
    output reg                          i_rd_w_src_wb,
    output reg                          i_csr_w_en_wb,
    output reg  [`CSR_ADDR_LEN-1:0]     i_csr_w_addr_wb,
    output reg  [ 2: 0]                 i_mem_r_ext_type_wb,
    output reg  [`BUSLEN-1:0]           i_regbus_C_wb,
    output reg  [`BUSLEN-1:0]           i_regbus_D_wb
);
    //---------- IF->ID ----------//
    always @(posedge clock) begin
        if(reset) begin
            i_pc_valid_id    <= 0;
            i_pc_id          <= 0;
            i_instruction_id <= 32'h13;
        end
        else if(flush) begin
            i_pc_valid_id    <= 0;
            i_pc_id          <= 0;
            i_instruction_id <= 32'h13;
        end
        else if(stall[0]) begin
            i_pc_valid_id    <= i_pc_valid_id;
            i_pc_id          <= i_pc_id;
            i_instruction_id <= i_instruction_id;
        end
        else begin
            i_pc_valid_id    <= o_pc_valid_if;
            i_pc_id          <= o_pc_if;
            i_instruction_id <= o_instruction_if;
        end
    end
    
    //---------- ID->EX ----------//
    always @(posedge clock) begin
        if(reset) begin
            i_regbus_A_ex         <= 0;
            i_regbus_B_ex         <= 0;
            i_imm_ex              <= 0;
            i_ALU_oprend_A_src_ex <= 0;
            i_ALU_oprend_B_src_ex <= 0;
            i_ALU_op_ex           <= 0;
            i_ALU_out_ext_type_ex <= 0;
            i_adder_src_ex        <= 0;
            i_pc_ex               <= 0;
            i_pc_valid_ex         <= 0;
            i_mem_r_en_ex         <= 0;
            i_mem_r_ext_type_ex   <= 0;
            i_mem_w_en_ex         <= 0;
            i_mem_w_size_ex       <= 0;
            i_inst_jump_ex        <= 0;
            i_inst_branch_type_ex <= 3'b010;
            i_is_system_inst_ex   <= 0;
            i_inst_ecall_ex       <= 0;
            i_inst_mret_ex        <= 0;
            i_rd_w_en_ex          <= 0;
            i_rd_w_src_ex         <= 0;
            i_rd_w_addr_ex        <= 0;
            i_csr_w_en_ex         <= 0;
            i_csr_w_addr_ex       <= 0;
        end
        else if(flush) begin
            i_regbus_A_ex         <= 0;
            i_regbus_B_ex         <= 0;
            i_imm_ex              <= 0;
            i_ALU_oprend_A_src_ex <= 0;
            i_ALU_oprend_B_src_ex <= 0;
            i_ALU_op_ex           <= 0;
            i_ALU_out_ext_type_ex <= 0;
            i_adder_src_ex        <= 0;
            i_pc_ex               <= 0;
            i_pc_valid_ex         <= 0;
            i_mem_r_en_ex         <= 0;
            i_mem_r_ext_type_ex   <= 0;
            i_mem_w_en_ex         <= 0;
            i_mem_w_size_ex       <= 0;
            i_inst_jump_ex        <= 0;
            i_inst_branch_type_ex <= 3'b010;
            i_is_system_inst_ex   <= 0;
            i_inst_ecall_ex       <= 0;
            i_inst_mret_ex        <= 0;
            i_rd_w_en_ex          <= 0;
            i_rd_w_src_ex         <= 0;
            i_rd_w_addr_ex        <= 0;
            i_csr_w_en_ex         <= 0;
            i_csr_w_addr_ex       <= 0;
        end
        else if(stall[1]) begin
            i_regbus_A_ex         <= i_regbus_A_ex;
            i_regbus_B_ex         <= i_regbus_B_ex;
            i_imm_ex              <= i_imm_ex;
            i_ALU_oprend_A_src_ex <= i_ALU_oprend_A_src_ex;
            i_ALU_oprend_B_src_ex <= i_ALU_oprend_B_src_ex;
            i_ALU_op_ex           <= i_ALU_op_ex;
            i_ALU_out_ext_type_ex <= i_ALU_out_ext_type_ex;
            i_adder_src_ex        <= i_adder_src_ex;
            i_pc_ex               <= i_pc_ex;
            i_pc_valid_ex         <= i_pc_valid_ex;
            i_mem_r_en_ex         <= i_mem_r_en_ex;
            i_mem_r_ext_type_ex   <= i_mem_r_ext_type_ex;
            i_mem_w_en_ex         <= i_mem_w_en_ex;
            i_mem_w_size_ex       <= i_mem_w_size_ex;
            i_inst_jump_ex        <= i_inst_jump_ex;
            i_inst_branch_type_ex <= i_inst_branch_type_ex;
            i_is_system_inst_ex   <= i_is_system_inst_ex;
            i_inst_ecall_ex       <= i_inst_ecall_ex;
            i_inst_mret_ex        <= i_inst_mret_ex;
            i_rd_w_en_ex          <= i_rd_w_en_ex;
            i_rd_w_src_ex         <= i_rd_w_src_ex;
            i_rd_w_addr_ex        <= i_rd_w_addr_ex;
            i_csr_w_en_ex         <= i_csr_w_en_ex;
            i_csr_w_addr_ex       <= i_csr_w_addr_ex;
        end
        else if(bubble) begin
            i_regbus_A_ex         <= 0;
            i_regbus_B_ex         <= 0;
            i_imm_ex              <= 0;
            i_ALU_oprend_A_src_ex <= 0;
            i_ALU_oprend_B_src_ex <= 0;
            i_ALU_op_ex           <= 0;
            i_ALU_out_ext_type_ex <= 0;
            i_adder_src_ex        <= 0;
            i_pc_ex               <= 0;
            i_pc_valid_ex         <= 0;
            i_mem_r_en_ex         <= 0;
            i_mem_r_ext_type_ex   <= 0;
            i_mem_w_en_ex         <= 0;
            i_mem_w_size_ex       <= 0;
            i_inst_jump_ex        <= 0;
            i_inst_branch_type_ex <= 3'b010;
            i_is_system_inst_ex   <= 0;
            i_inst_ecall_ex       <= 0;
            i_inst_mret_ex        <= 0;
            i_rd_w_en_ex          <= 0;
            i_rd_w_src_ex         <= 0;
            i_rd_w_addr_ex        <= 0;
            i_csr_w_en_ex         <= 0;
            i_csr_w_addr_ex       <= 0;
        end
        else begin
            i_regbus_A_ex         <= o_regbus_A_id;
            i_regbus_B_ex         <= o_regbus_B_id;
            i_imm_ex              <= o_imm_id;
            i_ALU_oprend_A_src_ex <= o_ALU_oprend_A_src_id;
            i_ALU_oprend_B_src_ex <= o_ALU_oprend_B_src_id;
            i_ALU_op_ex           <= o_ALU_op_id;
            i_ALU_out_ext_type_ex <= o_ALU_out_ext_type_id;
            i_adder_src_ex        <= o_adder_src_id;
            i_pc_ex               <= o_pc_id;
            i_pc_valid_ex         <= o_pc_valid_id;
            i_mem_r_en_ex         <= o_mem_r_en_id;
            i_mem_r_ext_type_ex   <= o_mem_r_ext_type_id;
            i_mem_w_en_ex         <= o_mem_w_en_id;
            i_mem_w_size_ex       <= o_mem_w_size_id;
            i_inst_jump_ex        <= o_inst_jump_id;
            i_inst_branch_type_ex <= o_inst_branch_type_id;
            i_is_system_inst_ex   <= o_is_system_inst_id;
            i_inst_ecall_ex       <= o_inst_ecall_id;
            i_inst_mret_ex        <= o_inst_mret_id;
            i_rd_w_en_ex          <= o_rd_w_en_id;
            i_rd_w_src_ex         <= o_rd_w_src_id;
            i_rd_w_addr_ex        <= o_rd_w_addr_id;
            i_csr_w_en_ex         <= o_csr_w_en_id;
            i_csr_w_addr_ex       <= o_csr_w_addr_id;
        end
    end
    
    //---------- EX->MEM ----------//
    always @(posedge clock) begin
        if(reset) begin
            i_branch_target_pc_mem <= 0;
            i_ALU_out_mem          <= 0;
            i_psw_flags_mem        <= 0;
            i_mem_r_en_mem         <= 0;
            i_mem_r_ext_type_mem   <= 0;
            i_mem_w_en_mem         <= 0;
            i_mem_w_size_mem       <= 0;
            i_regbus_B_mem         <= 0;
            i_rd_w_addr_mem        <= 0;
            i_rd_w_en_mem          <= 0;
            i_rd_w_src_mem         <= 0;
            i_csr_w_en_mem         <= 0;
            i_csr_w_addr_mem       <= 0;
            i_inst_jump_mem        <= 0;
            i_inst_branch_type_mem <= 3'b010;
            i_is_system_inst_mem   <= 0;
            i_inst_ecall_mem       <= 0;
            i_inst_mret_mem        <= 0;
        end
        else if(flush) begin
            i_branch_target_pc_mem <= 0;
            i_ALU_out_mem          <= 0;
            i_psw_flags_mem        <= 0;
            i_mem_r_en_mem         <= 0;
            i_mem_r_ext_type_mem   <= 0;
            i_mem_w_en_mem         <= 0;
            i_mem_w_size_mem       <= 0;
            i_regbus_B_mem         <= 0;
            i_rd_w_addr_mem        <= 0;
            i_rd_w_en_mem          <= 0;
            i_rd_w_src_mem         <= 0;
            i_csr_w_en_mem         <= 0;
            i_csr_w_addr_mem       <= 0;
            i_inst_jump_mem        <= 0;
            i_inst_branch_type_mem <= 3'b010;
            i_is_system_inst_mem   <= 0;
            i_inst_ecall_mem       <= 0;
            i_inst_mret_mem        <= 0;
        end
        else if(stall[2]) begin
            i_branch_target_pc_mem <= i_branch_target_pc_mem;
            i_ALU_out_mem          <= i_ALU_out_mem;
            i_psw_flags_mem        <= i_psw_flags_mem;
            i_mem_r_en_mem         <= i_mem_r_en_mem;
            i_mem_r_ext_type_mem   <= i_mem_r_ext_type_mem;
            i_mem_w_en_mem         <= i_mem_w_en_mem;
            i_mem_w_size_mem       <= i_mem_w_size_mem;
            i_regbus_B_mem         <= i_regbus_B_mem;
            i_rd_w_addr_mem        <= i_rd_w_addr_mem;
            i_rd_w_en_mem          <= i_rd_w_en_mem;
            i_rd_w_src_mem         <= i_rd_w_src_mem;
            i_csr_w_en_mem         <= i_csr_w_en_mem;
            i_csr_w_addr_mem       <= i_csr_w_addr_mem;
            i_inst_jump_mem        <= i_inst_jump_mem;
            i_inst_branch_type_mem <= i_inst_branch_type_mem;
            i_is_system_inst_mem   <= i_is_system_inst_mem;
            i_inst_ecall_mem       <= i_inst_ecall_mem;
            i_inst_mret_mem        <= i_inst_mret_mem;
        end
        else begin
            i_branch_target_pc_mem <= o_pc_ex;
            i_ALU_out_mem          <= o_ALU_out_ex;
            i_psw_flags_mem        <= o_psw_flags_ex;
            i_mem_r_en_mem         <= o_mem_r_en_ex;
            i_mem_r_ext_type_mem   <= o_mem_r_ext_type_ex;
            i_mem_w_en_mem         <= o_mem_w_en_ex;
            i_mem_w_size_mem       <= o_mem_w_size_ex;
            i_regbus_B_mem         <= o_regbus_B_ex;
            i_rd_w_addr_mem        <= o_rd_w_addr_ex;
            i_rd_w_en_mem          <= o_rd_w_en_ex;
            i_rd_w_src_mem         <= o_rd_w_src_ex;
            i_csr_w_en_mem         <= o_csr_w_en_ex;
            i_csr_w_addr_mem       <= o_csr_w_addr_ex;
            i_inst_jump_mem        <= o_inst_jump_ex;
            i_inst_branch_type_mem <= o_inst_branch_type_ex;
            i_is_system_inst_mem   <= o_is_system_inst_ex;
            i_inst_ecall_mem       <= o_inst_ecall_ex;
            i_inst_mret_mem        <= o_inst_mret_ex;
        end
    end

    //---------- MEM->WB ----------//
    always @(posedge clock) begin
        if(reset) begin
            i_rd_w_addr_wb         <= 0;
            i_rd_w_en_wb           <= 0;
            i_rd_w_src_wb          <= 0;
            i_csr_w_en_wb          <= 0;
            i_csr_w_addr_wb        <= 0;
            i_mem_r_ext_type_wb    <= 0;
            i_regbus_C_wb          <= 0;
            i_regbus_D_wb          <= 0;
        end
        else if(stall[3]) begin
            i_rd_w_addr_wb         <= i_rd_w_addr_wb;
            i_rd_w_en_wb           <= i_rd_w_en_wb;
            i_rd_w_src_wb          <= i_rd_w_src_wb;
            i_csr_w_en_wb          <= i_csr_w_en_wb;
            i_csr_w_addr_wb        <= i_csr_w_addr_wb;
            i_mem_r_ext_type_wb    <= i_mem_r_ext_type_wb;
            i_regbus_C_wb          <= i_regbus_C_wb;
            i_regbus_D_wb          <= i_regbus_D_wb;
        end
        else begin
            i_rd_w_addr_wb         <= o_rd_w_addr_mem;
            i_rd_w_en_wb           <= o_rd_w_en_mem;
            i_rd_w_src_wb          <= o_rd_w_src_mem;
            i_csr_w_en_wb          <= o_csr_w_en_mem;
            i_csr_w_addr_wb        <= o_csr_w_addr_mem;
            i_mem_r_ext_type_wb    <= o_mem_r_ext_type_mem;
            i_regbus_C_wb          <= o_regbus_C_mem;
            i_regbus_D_wb          <= o_regbus_D_mem;
        end
    end

endmodule